# 🖼️ Tensor可视化工具
import streamlit as st
import torch
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
import plotly.express as px
import matplotlib.font_manager as fm
import matplotlib

st.set_page_config(page_title="tensor_visualizer", page_icon="🖼️")
# 设置页面标题
st.title("🚀 PyTorch Tensor可视化工具")
st.caption("作者：何双新 ｜ 环境：Mac M1 + PyTorch")


# 侧边栏选项
st.sidebar.header("Tensor设置")
tensor_dim = st.sidebar.radio("选择Tensor维度", [0, 1, 2, 3, 4], index=2)

# 根据维度提供不同选项
if tensor_dim == 0:  # 标量
    scalar_value = st.sidebar.slider("标量值", -10.0, 10.0, 5.0, 0.1)
    
    st.header("0维Tensor (标量)")
    tensor = torch.tensor(scalar_value)
    st.code(f"tensor = torch.tensor({scalar_value})")
    st.write(f"值: {tensor.item()}")
    st.write(f"形状: {tensor.shape}")
    
    # 可视化
    st.write("可视化: 一个点")
    fig, ax = plt.subplots(figsize=(3, 3))
    ax.scatter([0], [0], s=100, c=[scalar_value], cmap='viridis')
    ax.set_xlim(-1, 1)
    ax.set_ylim(-1, 1)
    ax.set_xticks([])
    ax.set_yticks([])
    st.pyplot(fig)
    
elif tensor_dim == 1:  # 向量
    vector_size = st.sidebar.slider("向量大小", 2, 20, 10)
    vector_type = st.sidebar.selectbox("向量类型", ["随机", "线性", "正弦波"])
    
    st.header("1维Tensor (向量)")
    
    if vector_type == "随机":
        tensor = torch.rand(vector_size)
    elif vector_type == "线性":
        tensor = torch.linspace(0, 10, vector_size)
    else:  # 正弦波
        tensor = torch.sin(torch.linspace(0, 6.28, vector_size))
    
    st.code(f"tensor.shape = {tensor.shape}")
    st.write("Tensor值:")
    st.write(tensor)
    
    # 可视化
    st.write("可视化:")
    fig, ax = plt.subplots(figsize=(10, 4))
    ax.plot(tensor.numpy(), marker='o')
    ax.set_title("1维Tensor可视化")
    ax.set_xlabel("索引")
    ax.set_ylabel("值")
    ax.grid(True)
    st.pyplot(fig)
    
elif tensor_dim == 2:  # 矩阵
    rows = st.sidebar.slider("行数", 2, 10, 5)
    cols = st.sidebar.slider("列数", 2, 10, 5)
    tensor_type = st.sidebar.selectbox("矩阵类型", ["随机", "单位矩阵", "对角矩阵"])
    
    st.header("2维Tensor (矩阵)")
    
    if tensor_type == "随机":
        tensor = torch.rand(rows, cols)
    elif tensor_type == "单位矩阵":
        tensor = torch.eye(max(rows, cols))[:rows, :cols]
    else:  # 对角矩阵
        tensor = torch.diag(torch.linspace(1, min(rows, cols), min(rows, cols)))
        if rows > cols:
            tensor = torch.cat([tensor, torch.zeros(rows - cols, cols)], dim=0)
        elif cols > rows:
            tensor = torch.cat([tensor, torch.zeros(rows, cols - rows)], dim=1)
    
    st.code(f"tensor.shape = {tensor.shape}")
    st.write("Tensor值:")
    st.write(tensor)
    
    # 可视化为热力图
    st.write("可视化:")
    fig = px.imshow(tensor.numpy(), 
                    labels=dict(x="列", y="行", color="值"),
                    color_continuous_scale='viridis')
    fig.update_layout(width=600, height=500)
    st.plotly_chart(fig)
    
elif tensor_dim == 3:  # 3D Tensor
    depth = st.sidebar.slider("深度", 2, 5, 3)
    height = st.sidebar.slider("高度", 2, 10, 5)
    width = st.sidebar.slider("宽度", 2, 10, 5)
    
    st.header("3维Tensor")
    tensor = torch.rand(depth, height, width)
    
    st.code(f"tensor.shape = {tensor.shape}")
    
    # 展示每个深度层
    st.write("每个深度的切片可视化:")
    
    tabs = st.tabs([f"切片 {i}" for i in range(depth)])
    for i, tab in enumerate(tabs):
        with tab:
            fig = px.imshow(tensor[i].numpy(),
                           labels=dict(x="宽度", y="高度", color="值"),
                           color_continuous_scale='viridis')
            fig.update_layout(width=500, height=400)
            st.plotly_chart(fig)
    
    # 3D可视化
    st.write("3D可视化 (体素):")
    # 创建网格
    X, Y, Z = np.mgrid[0:depth, 0:height, 0:width]
    values = tensor.numpy().flatten()
    
    fig = go.Figure(data=go.Volume(
        x=X.flatten(),
        y=Y.flatten(),
        z=Z.flatten(),
        value=values,
        opacity=0.1,
        surface_count=15,
        colorscale='viridis'
    ))
    fig.update_layout(
        scene=dict(xaxis_title='深度', yaxis_title='高度', zaxis_title='宽度'),
        width=700, height=700
    )
    st.plotly_chart(fig)
    
elif tensor_dim == 4:  # 4D Tensor
    batch = st.sidebar.slider("批量大小", 1, 5, 2)
    channels = st.sidebar.slider("通道数", 1, 3, 3)
    height = st.sidebar.slider("高度", 4, 12, 8)
    width = st.sidebar.slider("宽度", 4, 12, 8)
    
    st.header("4维Tensor (批量图像)")
    tensor = torch.rand(batch, channels, height, width)
    
    st.code(f"tensor.shape = {tensor.shape}")
    st.write(f"这个Tensor可以表示{batch}张{channels}通道的{height}x{width}图像")
    
    # 可视化每个批次的图像
    batch_tabs = st.tabs([f"批次 {i}" for i in range(batch)])
    
    for b, batch_tab in enumerate(batch_tabs):
        with batch_tab:
            if channels == 3:
                # 针对RGB图像的特殊处理
                img = tensor[b].permute(1, 2, 0).numpy().clip(0, 1)

                # img = tensor[b].permute(1, 2, 0).numpy()  # 转换为HWC格式
                st.image(img, caption=f"批次 {b} 的RGB图像", use_container_width=True)

            else:
                # 展示每个通道
                channel_tabs = st.tabs([f"通道 {i}" for i in range(channels)])
                for c, channel_tab in enumerate(channel_tabs):
                    with channel_tab:
                        fig = px.imshow(tensor[b, c].numpy(),
                                       color_continuous_scale='viridis')
                        fig.update_layout(width=400, height=400)
                        st.plotly_chart(fig)

# 添加信息部分
st.sidebar.markdown("---")
st.sidebar.info("""
这个应用程序帮助您可视化不同维度的PyTorch Tensor。
- 0维：标量（一个点）
- 1维：向量（一条线）
- 2维：矩阵（一个平面）
- 3维：3D张量（一个立方体）
- 4维：4D张量（批量图像）
""")

st.caption("安徽智加数字科技有限公司 · 技术学习组出品 🚀")

